package trie

import (
	"errors"
	"regexp"
	"strings"
)

var charsReg = regexp.MustCompile(`^[a-zA-Z]+$`)

// Trie: 只支持26个英文字母的字典树
type Trie struct {
	root *node
}

func NewTrie() *Trie {
	return &Trie{
		root: newNode(0, false),
	}
}

type node struct {
	leaf     bool
	char     uint8
	children []*node
}

func newNode(char uint8, leaf bool) *node {
	return &node{
		leaf:     leaf,
		char:     char,
		children: make([]*node, 26),
	}
}

func (t Trie) ListWords() []string {
	words := make([]string, 0)
	for _, n := range t.root.children {
		if n != nil {
			bytes := make([]byte, 0)
			bytes = append(bytes, n.char)
			dfs(&words, n, bytes)
		}
	}
	return words
}

func dfs(words *[]string, node *node, buffer []byte) {
	for _, n := range node.children {
		if n != nil {
			dfs(words, n, append(buffer, n.char))
			if n.leaf {
				*words = append(*words, string(append(buffer, n.char)))
			}
		}
	}
}

func (t *Trie) AddWord(word string) error {
	if charsReg.Match([]byte(word)) {
		addWord(t.root, strings.ToLower(word))
		return nil
	}
	return errors.New("only support 26 letters of an alphabet")
}

func addWord(node *node, word string) {
	if len(word) == 0 {
		return
	}
	c := word[0]
	index := c - 'a'
	if node.children[index] == nil {
		node.children[index] = newNode(c, len(word) == 1)
	}
	addWord(node.children[index], word[1:])
}
